# 词嵌入维度d_model
d_model = 512
class PositionalEncoding(nn.Module):
    def __init__(self, d_model, dropout=0.1, max_len=5000):
        super(PositionalEncoding, self).__init__()
        self.dropout = nn.Dropout(p=dropout)
        pe = torch.zeros(max_len, d_model) #(5000,512)
        # pe: [max_len,d_model]
        # 初始化size为 max_len(文本最大长度)*embedding维度的全零矩阵
        # 来存放所有小于这个长度位置对应的positional embedding
        position=torch.arange(0,max_len,dtype=torch.float).unsqueeze(1)
        # position:[max_len,1] (5000,1)
        # 生成一个位置下标的tensor矩阵(每一行都是一个位置下标)
        """
        形式如：
        tensor([[0.],
                 [1.],
                 [2.],
                ...])
        """
        div_term = torch.exp(torch.arange(0,d_model,2).float() *(-math.log(10000.0)/d_model))
        # div_term:[d_model/2]
        # e^(-i*log 10000/d_model)=10000^(-i/d_model)
        # d_model为embedding_dimension
        # 这里幂运算多,使用exp和log来转换公式中pos 

        # 两个相乘的维度为[max_len,d_model/2]
        pe[:, 0::2] = torch.sin(position * div_term)
        pe[:, 1::2] = torch.cos(position * div_term)
        # 计算position encoding:  
        # 编码方式：[sin,cos,sin,cos,sin,cos,...]
        # pe的维度为[max_len,d_model]
        # 每一行的奇数偶数分别取sin和cos(position * div_term)里面的值
        # 根据公式,各位置在各embedding维度上的位置纹理值,存放到pe矩阵中
        pe = pe.unsqueeze(0)
        # 维度变成(1, max_len,d_model)  (1,5000,512)
        # pe = pe.unsqueeze(0)
        # 加1个维度,使得pe维度变为(1,max_len,d_model) 
        # 方便后续与一个batch的句子所有词的embedding批量相加
        self.register_buffer('my_pe', pe)
        # 放入buffer中,参数不会训练
        # 将pe矩阵以持久的buffer状态存下(不会作为要训练的参数)
    def forward(self, x):
        '''
        x: [batch_size,seq_length, d_model]  (2,5,512)
        '''
        x = x + self.my_pe[:, :x.size(1), :] 
        # pe:(1,5000,512)->(1,5,512)
        # 一个batch的句子所有词的Embedding+构建好的PositionalEmbeding
        # 选取和x一样维度的seq_length,将pe加到x上, 
        # 按照该批次的最大句子长度来取对应需要的那些PositionalEmbedding值
        return self.dropout(x)  
